use std::{fs::File, path::Path};

use safetensors::SafeTensors;
use serde::Serialize;
use serde_json::{json, Value};

use crate::{config::LlamaConfigJson, params::LLamaParams};

pub struct Llama<T: Serialize> {
    vocab: usize,           // vocab size
    n_layers: usize,        // number of layers
    n_q_h: usize,           // number of heads for q
    n_kv_h: usize,          // number of heads for k and v
    d: usize,               // dimension of hidden states
    dqkv: usize,            // length of a single q, k, or v vector
    di: usize,              // dimension of intermediate states
    eps: f32,               // epsilon for RMS normalization
    rope_theta: f32,        // rope theta for rope initialization
    max_seq_len: usize,     // maximum sequence length
    params: LLamaParams<T>, // trained weights of this model
    bos_token_id: u32,      // start token id
    eos_token_id: u32,      // end token id
}

impl Llama<f32> {
    pub fn from_safetensors(model_dir: impl AsRef<Path>) -> Self {
        let config = File::open(model_dir.as_ref().join("config.json")).unwrap();
        let config: LlamaConfigJson = serde_json::from_reader(config).unwrap();
        let model_file = std::fs::read(model_dir.as_ref().join("model.safetensors")).unwrap();
        let safetensor = SafeTensors::deserialize(&model_file).unwrap();
        let params = LLamaParams::from_safetensors(&safetensor, &config);

        Self {
            vocab: config.vocab_size,
            n_layers: config.num_hidden_layers,
            n_q_h: config.num_attention_heads,
            n_kv_h: config.num_key_value_heads,
            d: config.hidden_size,
            dqkv: config.hidden_size / config.num_attention_heads,
            di: config.intermediate_size,
            eps: config.rms_norm_eps,
            rope_theta: config.rope_theta,
            max_seq_len: config.max_position_embeddings,
            params: params,
            bos_token_id: config.bos_token_id,
            eos_token_id: config.eos_token_id,
        }
    }

    pub fn to_json(&self) -> serde_json::Result<Value> {
        let params_json = self.params.to_json()?;

        Ok(json!({
            "vocab": self.vocab,
            "n_layers": self.n_layers,
            "n_q_h": self.n_q_h,
            "n_kv_h": self.n_kv_h,
            "d": self.d,
            "dqkv": self.dqkv,
            "di": self.di,
            "eps": self.eps,
            "rope_theta": self.rope_theta,
            "max_seq_len": self.max_seq_len,
            "params": params_json,
            "bos_token_id": self.bos_token_id,
            "eos_token_id": self.eos_token_id,
        }))
    }
}
